from typing import List

import fastapi
from fastapi.responses import StreamingResponse
from loguru import logger

from .api import (
    ResponseCheckTaskOutput,
    ResponseClearTask,
    ResponseGetLogs,
    ResponseGetTaskStatus,
    ResponsePostFile,
    ResponsePostFileChunk,
    ResponsePostFileFinish,
    ResponsePutTask,
    ResponseSetTaskStatus,
    TaskFile,
    TaskManager,
    TaskSpec,
    TaskStatus,
    ResponseGetTaskFiles,
)

app = fastapi.FastAPI()


@app.get("/")
def read_root():
    return "RR_OK"


task_manager = TaskManager()


@app.put("/tasks")
def put_task(tasks: List[TaskSpec]) -> ResponsePutTask:
    # PUT tasks to the database
    logger.info(f"Received {len(tasks)} tasks")
    for task in tasks:
        task_manager.put_task(task)
    return ResponsePutTask(success=True, message="Task received")


@app.get("/ready_task/")
def get_task() -> TaskSpec:
    # GET tasks from the database
    task = task_manager.acquire_ready_task()
    if task is None:
        raise fastapi.HTTPException(status_code=404, detail="No more task")
    return task


@app.get("/task_status/{task_hash}")
def get_task_status(task_hash: str) -> ResponseGetTaskStatus:
    try:
        status = task_manager.get_task_status(task_hash)
    except Exception:
        logger.exception(f"Error getting task status for {task_hash}")
        return ResponseGetTaskStatus(success=False, status=TaskStatus.UNKNOWN)
    return ResponseGetTaskStatus(success=True, status=status)


@app.post("/task_status/{task_hash}/{status}")
def set_task_status(task_hash: str, status: TaskStatus) -> ResponseSetTaskStatus:
    task_manager.set_task_status(task_hash, status)
    return ResponseSetTaskStatus(success=True)


@app.get("/check_task_output/{task_hash}/{index}")
def check_task_output(task_hash: str, index: int) -> ResponseCheckTaskOutput:
    file_record = task_manager.get_result(task_hash, index)
    task_status = task_manager.get_task_status(task_hash)
    return ResponseCheckTaskOutput(
        ready=file_record is not None,
        task_status=task_status,
        record=file_record,
    )


@app.get("/get_logs/{task_hash}/{limit}/{offset}")
def get_logs(task_hash: str, limit: int = 100, offset: int = 0) -> ResponseGetLogs:
    logs = list(task_manager.get_logs(task_hash, limit, offset))
    return ResponseGetLogs(success=True, logs=logs)


@app.post("/file")
def put_file(file: TaskFile) -> ResponsePostFile:
    task_manager.add_file(file)
    return ResponsePostFile(success=True)


@app.post("/file_chunk/{file_id}/{index}")
async def put_file_chunk(
    file_id: str,
    index: int,
    request: fastapi.Request,
) -> ResponsePostFileChunk:
    chunk = await request.body()
    task_manager.add_file_chunk(file_id, chunk, index)
    return ResponsePostFileChunk(success=True)


@app.post("/file_finish/{file_id}")
def finish_file(file_id: str) -> ResponsePostFileFinish:
    task_manager.finish_file(file_id)
    return ResponsePostFileFinish(success=True)


@app.get("/file/{file_id}")
def get_file(file_id: str) -> StreamingResponse:
    file = task_manager.get_file(file_id)
    file_info = task_manager.get_file_info(file_id)
    return StreamingResponse(file, media_type=file_info.content_type)


@app.get("/file_info/{file_id}")
def get_file_info(file_id: str) -> TaskFile:
    file_info = task_manager.get_file_info(file_id)
    if file_info is None:
        raise fastapi.HTTPException(status_code=404, detail="File not found")
    return file_info


@app.delete("/task/{task_hash}")
def clear_task(task_hash: str) -> ResponseClearTask:
    task_manager.clear_task(task_hash)
    return ResponseClearTask(success=True)


@app.get("/task_spec/{hash}")
def get_task_spec(hash: str) -> TaskSpec:
    task_spec = task_manager._get_by_hash(hash)
    if task_spec is None:
        raise fastapi.HTTPException(status_code=404, detail="Task not found")
    return task_spec


@app.get("/task_files/{task_hash}")
def get_task_files(task_hash: str) -> ResponseGetTaskFiles:
    files = task_manager.get_task_files(task_hash)
    return ResponseGetTaskFiles(success=True, files=files)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000, root_path="/rr_server")
